﻿// onnx_test.cpp: 定义应用程序的入口点。
//

#include "onnx_test.h"

#include <iostream>
#include <assert.h>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include <string>
#include <chrono>


bool input_config(bool &boolean)
{
    std::string instr;
    std::cin >> instr;
    if (instr[0] == 'y' || instr[0] == 'Y')
        boolean = true;
    return boolean;
}

int main(int argc, char* argv[]) {

    const wchar_t* model_path = L"test.onnx";
    const wchar_t* optimized_file_path = L"optimized.onnx";

    //Ort::Session session{ env, model_path, session_options };
	// print model input layer (node names, types, shape etc.)
	Ort::AllocatorWithDefaultOptions allocator;

    bool VERBOSE = false;
    bool USEGPU = false;

    try {

        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
        Ort::SessionOptions session_options;


        session_options.SetIntraOpNumThreads(4);
        //session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
        session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
        session_options.SetOptimizedModelFilePath(optimized_file_path);

        OrtCUDAProviderOptions gpu_options;
        gpu_options.device_id = 0;
        gpu_options.arena_extend_strategy = 0;
        gpu_options.gpu_mem_limit = (size_t)1 * 1024 * 1024 * 1024; //onnxruntime1.8.1, onnxruntime1.9.0
        gpu_options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::OrtCudnnConvAlgoSearchExhaustive;
        gpu_options.do_copy_in_default_stream = 1;
        



        std::cout << "Verbose y/n ?"<<std::endl;
        input_config(VERBOSE);

        // 获取输入节点信息

        std::cout << "Use GPU y/n ?" << std::endl;
        if (input_config(USEGPU))
        {
            session_options.AppendExecutionProvider_CUDA(gpu_options);
        }
            

        Ort::Session session(env, model_path, session_options);



        if (VERBOSE) std::cout << "[Info]: Loaded Model." << std::endl;
        size_t num_input_nodes = session.GetInputCount();
        size_t num_output_nodes = session.GetOutputCount();


        // 定义输入和输出节点的名称向量
        std::vector<const char*> input_node_names;
        std::vector<const char*> output_node_names;

        // 获取输入节点信息并填充到向量中

        Ort::AllocatedStringPtr in_name = session.GetInputNameAllocated(0, allocator);
        const char* in_name_cstr = in_name.get(); // 获取字符串指针
        if (VERBOSE)
            std::cout << "[Info]: Input Name: " << in_name_cstr << std::endl;
        input_node_names.push_back(in_name_cstr);


    // 获取输出节点信息并填充到向量中

        Ort::AllocatedStringPtr out_name = session.GetOutputNameAllocated(0, allocator);
        const char* out_name_cstr = out_name.get(); // 获取字符串指针
        if (VERBOSE)
            std::cout << "[Info]: Output Name: " << out_name_cstr << std::endl;
        output_node_names.push_back(out_name_cstr);



        // 设置输入数据的维度，这里以单条数据为例
        std::vector<int64_t> input_node_dims = { 1, 1 };

        size_t input_tensor_size = 1*1;
        // 构造输入数据

        float inin;
        std::cout << "Input the phase：";
        std::cin >> inin;

        auto start_time = std::chrono::system_clock::now();

        std::vector<float> input_tensor_values(input_tensor_size);

        input_tensor_values[0] = inin;
        std::cout << input_tensor_values[0] << std::endl;

        // create input tensor object from data values
        auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
            memory_info,
            input_tensor_values.data(),
            input_tensor_size,
            input_node_dims.data(),
            input_node_dims.size()
            );
        assert(input_tensor.IsTensor());
        //std:: cout << input_node_names.data() << std::endl;
        std::vector<Ort::Value> ort_inputs;
        ort_inputs.push_back(std::move(input_tensor));
        // score model & input tensor, get back output tensor
        auto start_time1 = std::chrono::system_clock::now();
        auto output_tensors = session.Run(
            Ort::RunOptions{ nullptr },
            input_node_names.data(),
            ort_inputs.data(),
            ort_inputs.size(),
            output_node_names.data(),
            1
        );
        std::chrono::duration<double> diff1 = std::chrono::system_clock::now() - start_time1;


        if (VERBOSE) std::cout << "[Info]: Infer Cost time=" << diff1.count() * 1000.0 << "ms" << std::endl;
        // Get pointer to output tensor float values
        //float* floatarr = output_tensors[0].GetTensorMutableData<float>();



        Ort::Value& output_tensor = output_tensors[0];
        const float* output_data = output_tensor.GetTensorData<float>();
        std::vector<std::vector<float>> out_mat;
        for (int i = 0; i < 129; i++)
        {
            std::vector<float>line;
            for (int j = 0; j < 129; j++)
            {
                line.push_back(-output_data[i * 129 + j]);
            }
            out_mat.push_back(line);
        
        }        

        std::chrono::duration<double> diff = std::chrono::system_clock::now() - start_time;
        if (VERBOSE) std::cout << "[Info]: Whole Cost time=" << diff.count()*1000.0 << "ms" << std::endl;
        std::cout << "Show result? y/n" << std::endl;



        //是否打印显示
        bool SHOW = false;

        if (input_config(SHOW))
        {
            for (int i = 0; i < 129; i++)
            {
                std::vector<float>line;
                for (int j = 0; j < 129; j++)
                {
                    std::cout << " " << -output_data[i * 129 + j];
                }
                std::cout << std::endl;

            }
        }
        

        
    }

    catch (const Ort::Exception& e) {
        // 处理 Ort::Exception 异常
        std::cerr << "Caught Ort::Exception: " << std::string(e.what()) << std::endl;

        // 在异常描述信息中查找错误代码
        size_t pos = std::string(e.what()).find("ErrorCode: ");
        if (pos != std::string::npos) {
            std::string error_code_str = std::string(e.what()).substr(pos + 12); // 12 是 "ErrorCode: " 的长度
            int error_code = std::stoi(error_code_str);
            std::cerr << "Error Code: " << error_code << std::endl;
        }

        // 可选：进行其他异常处理或返回错误码
        return -1;
    }
    system("pause");
	return 0;
}